#include "RayTracer.h"

#include "KdTreeMemoryManager.h"
#include "TriangleListMemoryManager.h"
#include "Intersections.h"
#include "KdTraversalStack.h"
#include "LightSource.h"
#include "ShadowRay.h"
#include "MinMaxBin.h"
#include "KdBuildQueue.h"
#include "MinMax.h"
#include <assert.h>

extern BITMAP* buffer;
extern KdTreeMemoryManager kdtmm;
extern TriangleListMemoryManager tlmm;
extern SPoint viewPoint;
extern KdTraversalStack stack;
extern KdBuildQueue queue;
extern MinMaxBin bin;

static float calculateCostFast(const AABB& aabb, const float splitpos,
		const int axis, const int numLeft, const int numRight)
{
	const float size = .5 / aabb.getHullSize();

	return TRAVERSALCOST + INTERSECTCOST
			* (aabb.getLeftHullSize(splitpos, axis) * size * numLeft
					+ aabb.getRightHullSize(splitpos, axis) * size * numRight);
}

//static float calculateCost(Triangle** triangles, const int numTriangles,
//		const AABB& aabb, const float splitpos, const int axis, int& numLeft,
//		int&numRight)
//{
//	const float leftSize = aabb.getLeftHullSize(splitpos, axis);
//	const float rightSize = aabb.getRightHullSize(splitpos, axis);
//	const float size = .5 / aabb.getHullSize();
//
//	numLeft = numRight = 0;
//	for (int i = 0; i < numTriangles; i++)
//	{
//		if (triangles[i]->getMin(axis) < splitpos)
//			numLeft++;
//		if (triangles[i]->getMax(axis) >= splitpos)
//			numRight++;
//	}
//
//	return TRAVERSALCOST + INTERSECTCOST * (leftSize * size * numLeft
//			+ rightSize * size * numRight);
//}

#ifdef USERECURSIVEBUILD

static void _recBuildKdTree(KdNode* n, Triangle** tl, const int numTriangles,
		const AABB& aabb, int depth)
{
	// stop-criterion part one
	if (numTriangles <= MAXTRIANGLESPERKDNODE || depth >= MAXKDTREEDEPTH)
	{
		// add all triangles to this node
		Triangle** t = new (tlmm.request((numTriangles + 1))) Triangle*[numTriangles + 1];
		for (int i = 0; i < numTriangles; i++)
			*(t + i) = *(tl + i);
		*(t + numTriangles) = 0;

		n->setTriangleList(t);
		n->setLeaf(true);

		return;
	}

	// set the axis of the node
	const int axis = aabb.getMaxAxis();
	n->setAxis(axis);

	// setup the bin
	bin.clear();

	int numLeft = 0;
	int numRight = 0;
	float split = MAXFLOAT;
	float bestcost = MAXFLOAT;
	int tmpNumLeft = 0;
	int tmpNumRight = 0;
	//	if (1)//numTriangles > NUMBINS)
	//	{
	// determine the bin-size
	const float inverseBinSize = NUMBINS / aabb.getSize(axis);
	const float binSize = aabb.getSize(axis) / NUMBINS;

	// pass over triangles and bin them
	for (int i = 0; i < numTriangles; i++)
	{
		const Triangle* triangle = tl[i];
		const float leftExtreme = triangle->getMin(axis);
		const float rightExtreme = triangle->getMax(axis);
		bin.addMin(sMin(NUMBINS - 1, sMax(0, (int)((leftExtreme - aabb.getMin().get(axis)) * inverseBinSize))));
		bin.addMax(sMin(NUMBINS - 1, sMax(0, (int)((rightExtreme - aabb.getMin().get(axis)) * inverseBinSize))));
	}

	for (int i = 1; i < NUMBINS; i++)
		tmpNumRight += bin.getMax(i);
	tmpNumLeft = bin.getMin(0);

	// pass over bins to find split
	for (int i = 1; i < NUMBINS; i++)
	{
		const float tmpSplit = aabb.getMin().get(axis) + i * binSize;

		const float cost = calculateCostFast(aabb, tmpSplit, axis, tmpNumLeft,
				tmpNumRight);
		if (cost < bestcost)
		{
			bestcost = cost;
			split = tmpSplit;
			numLeft = tmpNumLeft;
			numRight = tmpNumRight;
		}

		tmpNumLeft += bin.getMin(i);
		tmpNumRight -= bin.getMax(i);
	}
	//	}
	//	else
	//	{
	//		// find the split-plane and the number of triangles on each side
	//		for (int i = 0; i < numTriangles; i++)
	//		{
	//			const Triangle* triangle = tl[i];
	//			const float leftExtreme = triangle->getMin(axis);
	//			const float rightExtreme = triangle->getMax(axis);
	//			if (leftExtreme > aabb.getMin().get(axis))
	//			{
	//				float cost = calculateCost(tl, numTriangles, aabb, leftExtreme,
	//						axis, tmpNumLeft, tmpNumRight);
	//				if (cost < bestcost)
	//				{
	//					bestcost = cost;
	//					split = leftExtreme;
	//					numLeft = tmpNumLeft;
	//					numRight = tmpNumRight;
	//				}
	//			}
	//
	//			if (rightExtreme < aabb.getMax().get(axis))
	//			{
	//				float cost = calculateCost(tl, numTriangles, aabb,
	//						rightExtreme, axis, tmpNumLeft, tmpNumRight);
	//				if (cost < bestcost)
	//				{
	//					bestcost = cost;
	//					split = rightExtreme;
	//					numLeft = tmpNumLeft;
	//					numRight = tmpNumRight;
	//				}
	//			}
	//		}
	//	}

	// stop-criterion part two
	//	if (bestcost >= numTriangles * INTERSECTCOST)
	//	{
	//		// add all triangles to this node
	//		Triangle** t = new (tlmm.request((numTriangles + 1))) Triangle*[numTriangles + 1];
	//		for (int i = 0; i < numTriangles; i++)
	//			*(t + i) = *(tl + i);
	//		*(t + numTriangles) = 0;
	//
	//		n->setTriangleList(t);
	//		n->setLeaf(true);
	//
	//		return;
	//	}

	n->setLeaf(false);

	// set the split-pos
	n->setSplitPos(split);

	//	std::cout << split << std::endl;

	Triangle** left = new (tlmm.request(numLeft)) Triangle*[numLeft];
	int leftCount = 0;
	Triangle** right = new (tlmm.request(numRight)) Triangle*[numRight];
	int rightCount = 0;
	for (int i = 0; i < numTriangles; i++)
	{
		if (tl[i]->getMin(axis) <= split)
			left[leftCount++] = tl[i];
		if (tl[i]->getMax(axis) > split)
			right[rightCount++] = tl[i];
	}

	assert(leftCount == numLeft);
	assert(rightCount == numRight);

	// request a place for the children-nodes
	KdNode* children = new (kdtmm.request(2)) KdNode[2];

	// link the current node to it's children
	n->setLeft(children);

	const AABB leftAABB = aabb.createLeft(split, axis);
	const AABB rightAABB = aabb.createRight(split, axis);
	_recBuildKdTree(children, left, numLeft, leftAABB, depth + 1);
	_recBuildKdTree(children + 1, right, numRight, rightAABB, depth + 1);
}

KdNode* recBuildKdTree(const vector<Triangle*>& triangles, const AABB& aabb)
{
	KdNode* root = new (kdtmm.request(1)) KdNode;

	// create the triangle-list
	const int numTriangles = triangles.size();
	Triangle** tl = new (tlmm.request(numTriangles)) Triangle*[numTriangles];

	// store the triangles in the list
	int index = 0;
	for (vector<Triangle*>::const_iterator it = triangles.begin(); it
			!= triangles.end(); it++)
	{
		tl[index] = *it;
		index++;
	}

	_recBuildKdTree(root, tl, numTriangles, aabb, 0);
	return root;
}

#else

KdNode* itBuildKdTree(const vector<Triangle*>& triangles, const AABB& aabb)
{
	const int numNodes = 2<<MAXKDTREEDEPTH;
	KdNode* root = new (kdtmm.request(numNodes)) KdNode[numNodes];
	KdNode* next = root + 1;

	// create the triangle-list
	const int numTriangles = triangles.size();
	Triangle** tl = new (tlmm.request(numTriangles)) Triangle*[numTriangles];

	// store the triangles in the list
	int index = 0;
	for (vector<Triangle*>::const_iterator it = triangles.begin(); it
			!= triangles.end(); it++)
	{
		tl[index] = *it;
		index++;
	}

	// clear the queue
	queue.clear();

	// add the root-node
	queue.add(root, tl, numTriangles, new AABB(aabb), 0);

	while (!queue.isEmpty())
	{ // while the queue still has items

		// get the first item in the queue
		QueueEntry& entry = queue.remove();

		// stop-criterion
		if (entry.numTriangles <= MAXTRIANGLESPERKDNODE || entry.depth
				>= MAXKDTREEDEPTH)
		{
			// add all triangles to this node
			Triangle
			** t =
			new (tlmm.request((entry.numTriangles + 1))) Triangle*[entry.numTriangles + 1];
			for (int i = 0; i < entry.numTriangles; i++)
			*(t + i) = *(entry.triangleList + i);
			*(t + entry.numTriangles) = 0;

			entry.node->setTriangleList(t);
			entry.node->setLeaf(true);

			delete entry.aabb;

			continue;
		}

		// set the axis of the node
		const int axis = entry.aabb->getMaxAxis();
		entry.node->setAxis(axis);

		// setup the bin
		bin.clear();

		int numLeft = 0;
		int numRight = 0;
		float split = MAXFLOAT;
		float bestcost = MAXFLOAT;
		int tmpNumLeft = 0;
		int tmpNumRight = 0;

		// determine the bin-size
		const float inverseBinSize = NUMBINS / entry.aabb->getSize(axis);
		const float binSize = entry.aabb->getSize(axis) / NUMBINS;

		// pass over triangles and bin them
		for (int i = 0; i < entry.numTriangles; i++)
		{
			const Triangle* triangle = entry.triangleList[i];
			const float leftExtreme = triangle->getMin(axis);
			const float rightExtreme = triangle->getMax(axis);
			bin.addMin(sMin(NUMBINS - 1, sMax(0, (int)((leftExtreme
													- entry.aabb->getMin().get(axis)) * inverseBinSize))));
			bin.addMax(sMin(NUMBINS - 1, sMax(0, (int)((rightExtreme
													- entry.aabb->getMin().get(axis)) * inverseBinSize))));
		}

		for (int i = 1; i < NUMBINS; i++)
		tmpNumRight += bin.getMax(i);
		tmpNumLeft = bin.getMin(0);

		// pass over bins to find split
		for (int i = 1; i < NUMBINS; i++)
		{
			const float tmpSplit = entry.aabb->getMin().get(axis) + i * binSize;

			const float cost = calculateCostFast(*(entry.aabb), tmpSplit, axis,
					tmpNumLeft, tmpNumRight);
			if (cost < bestcost)
			{
				bestcost = cost;
				split = tmpSplit;
				numLeft = tmpNumLeft;
				numRight = tmpNumRight;
			}

			tmpNumLeft += bin.getMin(i);
			tmpNumRight -= bin.getMax(i);
		}

		// set the node as a non-leaf
		entry.node->setLeaf(false);

		// set the split-pos
		entry.node->setSplitPos(split);

		Triangle** left = new (tlmm.request(numLeft)) Triangle*[numLeft];
		int leftCount = 0;
		Triangle** right = new (tlmm.request(numRight)) Triangle*[numRight];
		int rightCount = 0;
		for (int i = 0; i < entry.numTriangles; i++)
		{
			if (entry.triangleList[i]->getMin(axis) < split)
			left[leftCount++] = entry.triangleList[i];
			if (entry.triangleList[i]->getMax(axis) >= split)
			right[rightCount++] = entry.triangleList[i];
		}

		// link the current node to it's children
		entry.node->setLeft(next);

		AABB* leftAABB = new AABB(entry.aabb->createLeft(split, axis));
		AABB* rightAABB = new AABB(entry.aabb->createRight(split, axis));

		// delete aabb
		delete entry.aabb;

		// add both children to the queue
		// left
		queue.add(next, left, numLeft, leftAABB, entry.depth + 1);
		// right
		queue.add(next + 1, right, numRight, rightAABB, entry.depth + 1);

		next += 2;
	}

	return root;
}

#endif

void printKdTree(const KdNode* n)
{
	std::cout << "Address: " << n << std::endl;
	std::cout << "Is leaf: " << n->isLeaf() << std::endl;
	if (!n->isLeaf())
	{
		std::cout << "Split-position: " << n->getSplitPos() << std::endl;
		std::cout << "Split-axis: " << n->getAxis() << std::endl;
		std::cout << "Address of left child: " << n->getLeft() << std::endl;
		std::cout << "Address of right child: " << n->getRight() << std::endl;
		printKdTree(n->getLeft());
		printKdTree(n->getRight());
	}
	else
	{
		Triangle** cur = n->getTriangleList();
		int num = 0;
		while (*cur != 0)
		{
			num++;
			cur++;
		}
		std::cout << "Address of first triangle: " << n->getTriangleList()
				<< std::endl;
		std::cout << "Number of triangles: " << num << std::endl;
	}
}

static void _drawKdTree(const KdNode* n, int omitAxis, int x, int y,
		const AABB& aabb)
{
	int axis = n->getAxis();

	if (!n->isLeaf())
	{
		float split = n->getSplitPos();

		//		std::cout << split << std::endl;
		//		std::cout << min << std::endl;
		//		std::cout << max << std::endl;
		//		std::cout << std::endl;

		if (omitAxis == X)
		{
			if (axis == Y)
			{
				float min = aabb.getMin().get(Z);
				float max = aabb.getMax().get(Z);
				line(buffer, x + (int)(min * 10), y + (int)(split * 10), x
						+ (int)(max * 10), y + (int)(split * 10), makecol(0,
						255, 0));
			}
			else if (axis == Z)
			{
				float min = aabb.getMin().get(Y);
				float max = aabb.getMax().get(Y);
				line(buffer, x + (int)(split * 10), y + (int)(min * 10), x
						+ (int)(split * 10), y + (int)(max * 10), makecol(0, 0,
						255));
			}
		}
		else if (omitAxis == Y)
		{
			if (axis == X)
			{
				float min = aabb.getMin().get(Z);
				float max = aabb.getMax().get(Z);
				line(buffer, x + (int)(split * 10), y + (int)(min * 10), x
						+ (int)(split * 10), y + (int)(max * 10), makecol(255,
						0, 0));
			}
			else if (axis == Z)
			{
				float min = aabb.getMin().get(X);
				float max = aabb.getMax().get(X);
				line(buffer, x + (int)(min * 10), y + (int)(split * 10), x
						+ (int)(max * 10), y + (int)(split * 10), makecol(0, 0,
						255));
			}
		}
		else if (omitAxis == Z)
		{
			if (axis == X)
			{
				float min = aabb.getMin().get(Y);
				float max = aabb.getMax().get(Y);
				line(buffer, x + (int)(split * 10), y + (int)(min * 10), x
						+ (int)(split * 10), y + (int)(max * 10), makecol(255,
						0, 0));
			}
			else if (axis == Y)
			{
				float min = aabb.getMin().get(X);
				float max = aabb.getMax().get(X);
				line(buffer, x + (int)(min * 10), y + (int)(split * 10), x
						+ (int)(max * 10), y + (int)(split * 10), makecol(0,
						255, 0));
			}
		}

		const AABB leftAABB = aabb.createLeft(split, axis);
		const AABB rightAABB = aabb.createRight(split, axis);
		_drawKdTree(n->getLeft(), omitAxis, x, y, leftAABB);
		_drawKdTree(n->getRight(), omitAxis, x, y, rightAABB);
	}
}

void drawKdTree(const KdNode* n, int omitAxis, int x, int y, const AABB& aabb)
{
	if (omitAxis == X)
		rect(buffer, x + (int)(aabb.getMin().get(Z) * 10), y + (int)(aabb.getMin().get(Y) * 10), x + (int)(aabb.getMax().get(Z) * 10), y + (int)(aabb.getMax().get(Y) * 10), makecol(255, 255, 255));
	else if (omitAxis == Y)
		rect(buffer, x + (int)(aabb.getMin().get(X) * 10), y + (int)(aabb.getMin().get(Z) * 10), x + (int)(aabb.getMax().get(X) * 10), y + (int)(aabb.getMax().get(Z) * 10), makecol(255, 255, 255));
	else if (omitAxis == Z)
		rect(buffer, x + (int)(aabb.getMin().get(X) * 10), y + (int)(aabb.getMin().get(Y) * 10), x + (int)(aabb.getMax().get(X) * 10), y + (int)(aabb.getMax().get(Y) * 10), makecol(255, 255, 255));
	_drawKdTree(n, omitAxis, x, y, aabb);
}

/**
 * is the ray intersected?
 */
static bool intersect(const KdNode* node, ShadowRay& shadowRay)
{
	float bestDistance = MAXFLOAT;
	Triangle** cur = node->getTriangleList();
	while (*cur != 0)
	{
		//			distance = Intersections::rayTriangle(ray.getPoint(),
		//					ray.getVector(), (*cur)->getPoint(0), (*cur)->getPoint(1), (*cur)->getPoint(2));
		bestDistance = sMin(Intersections::rayTriangleFast(
				shadowRay.getStart(), shadowRay.getVector(), (*cur)->triAccel), bestDistance);
		cur++;
	}

	return bestDistance <= shadowRay.getLength();
}
/**
 * is the ray intersected?
 */
static void intersect(const KdNode* node, Ray& ray)
{
	float distance;
	Triangle** cur = node->getTriangleList();
	while (*cur != 0)
	{
		//			distance = Intersections::rayTriangle(ray.getPoint(),
		//					ray.getVector(), (*cur)->getPoint(0), (*cur)->getPoint(1), (*cur)->getPoint(2));
		distance = Intersections::rayTriangleFast(ray.getPoint(),
				ray.getVector(), (*cur)->triAccel);
		if (distance < ray.getHitDistance())
		{
			ray.setHitDistance(distance);
			ray.setHitTriangle(*cur);
		}
		cur++;
	}
}

#ifdef USERECURSIVETRAVERSAL

static float _recTraverse(KdNode* node, Ray& ray, float t_near, float t_far)
// returns distance to closest hit point

{
	if (node->isLeaf())
	{
		intersect(node, ray);
		return ray.getHitDistance();
	}

	// determine the front- and back-child
	KdNode* front;
	KdNode* back;
	if (viewPoint.get(node->getAxis()) > node->getSplitPos())
	{
		front = node->getRight();
		back = node->getLeft();
	}
	else
	{
		front = node->getLeft();
		back = node->getRight();
	}

	float d = (node->getSplitPos() - ray.getPoint().get(node->getAxis())) * ray.getInverseVector().get(node->getAxis());
	
	assert(d != std::numeric_limits<float>::infinity()); // check for infinity
	assert(d == d); // check for NaN

	if (d >= t_far || d < .0)
		return _recTraverse(front, ray, t_near, t_far);
	else if (d <= t_near)
		return _recTraverse(back, ray, t_near, t_far);
	else
	{
		// case three: traverse both sides in turn
		float distance = _recTraverse(front, ray, t_near, d);

		return (distance <= d) ? distance : _recTraverse(back, ray, d, t_far);
	}
}

float recTraverse(KdNode* root, Ray& ray, const AABB& aabb)
{
	float t_near = -1000;
	float t_far = 100000.;
	ray.setHitTriangle(0);
	ray.setHitDistance(MAXFLOAT);

	return (Intersections::AABBRay(t_near, t_far, aabb.getMin(), aabb.getMax(),
						ray.getPoint(), ray.getInverseVector())) ? _recTraverse(root, ray,
				t_near, t_far) : MAXFLOAT;
}

static bool _recTraverseShadow(KdNode* node, ShadowRay& shadowRay,
		float t_near, float t_far)
// returns distance to closest hit point

{
	if (t_near >= shadowRay.getLength())
		return true;

	if (node->isLeaf())
		return !intersect(node, shadowRay);

	// determine the front- and back-child
	KdNode* front;
	KdNode* back;
	if (shadowRay.getStart().get(node->getAxis()) > node->getSplitPos())
	{
		front = node->getRight();
		back = node->getLeft();
	}
	else
	{
		front = node->getLeft();
		back = node->getRight();
	}

	float d = (node->getSplitPos() - shadowRay.getStart().get(node->getAxis())) * shadowRay.getInverseVector().get(node->getAxis());

	assert(d != std::numeric_limits<float>::infinity()); // check for infinity
	assert(d == d); // check for NaN

	if (d >= t_far || d < .0)
		return _recTraverseShadow(front, shadowRay, t_near, t_far);
	else if (d <= t_near)
		return _recTraverseShadow(back, shadowRay, t_near, t_far);
	else
	{
		// case three: traverse both sides in turn
		if (!_recTraverseShadow(front, shadowRay, t_near, d))
			return false;

		return _recTraverseShadow(back, shadowRay, d, t_far);
	}
}

/**
 * does the shadow-ray reach the target?
 */
bool recTraverseShadow(KdNode* root, ShadowRay& shadowRay, const AABB& aabb)
{
	float t_near = 0;
	float t_far = shadowRay.getLength();

	return _recTraverseShadow(root, shadowRay, t_near, t_far);
}

#else

static float _itTraverseShadow(const KdNode* node, ShadowRay& shadowRay,
		float t_near, float t_far)
{
	while (1)
	{
		while (!node->isLeaf())
		{
			// determine the front- and back-child
			KdNode* front;
			KdNode* back;
			if (shadowRay.getStart().get(node->getAxis())> node->getSplitPos())
			{
				front = node->getRight();
				back = node->getLeft();
			}
			else
			{
				front = node->getLeft();
				back = node->getRight();
			}

			// traverse ’til next leaf
			float d = (node->getSplitPos() - shadowRay.getStart().get(node->getAxis())) * shadowRay.getInverseVector().get(node->getAxis());

			assert(d != std::numeric_limits<float>::infinity());
			// check for infinity
			assert(d == d);
			// check for NaN

			if (d >= t_far || d < .0)
			node = front;
			else if (d <= t_near)
			node = back;
			else
			{
				// case three: traverse both sides in turn
				stack.push(back, d, t_far);
				node = front;
				t_far = d;
			}
		}

		if (intersect(node, shadowRay))
		return false;
		if (stack.isEmpty())
		return true; // noting else to traverse any more...
		StackEntry s = stack.pop();
		node = s.node;
		t_near = s.near;
		t_far = s.far;
	}
}

bool itTraverseShadow(KdNode* root, ShadowRay& shadowRay, const AABB& aabb)
{
	float t_near = .0;
	float t_far = shadowRay.getLength();

	stack.clear();

	return _itTraverseShadow(root, shadowRay, t_near, t_far);
}

static float _itTraverse(const KdNode* node, Ray& ray, float t_near, float t_far)
{
	float d;
	while (1)
	{
		while (!node->isLeaf())
		{
			// determine the front- and back-child
			KdNode* front;
			KdNode* back;
			if (viewPoint.get(node->getAxis()) > node->getSplitPos())
			{
				front = node->getRight();
				back = node->getLeft();
			}
			else
			{
				front = node->getLeft();
				back = node->getRight();
			}

			// traverse ’til next leaf
			d = (node->getSplitPos() - ray.getPoint().get(node->getAxis())) * ray.getInverseVector().get(node->getAxis());
			
			assert(d != std::numeric_limits<float>::infinity()); // check for infinity
			assert(d == d); // check for NaN
			
			if (d >= t_far || d < .0)
			node = front;
			else if (d <= t_near)
			node = back;
			else
			{
				// case three: traverse both sides in turn
				stack.push(back, d, t_far);
				node = front;
				t_far = d;
			}
		}
		// have a leaf now
		intersect(node, ray);
		if (ray.getHitTriangle() && ray.getHitDistance() < t_far)
		return ray.getHitDistance(); // early ray termination
		if (stack.isEmpty())
		return ray.getHitDistance(); // noting else to traverse any more...
		StackEntry s = stack.pop();
		node = s.node;
		t_near = s.near;
		t_far = s.far;
	}
}

float itTraverse(KdNode* root, Ray& ray, const AABB& aabb)
{
	float t_near = -1000;
	float t_far = 100000.;
	ray.setHitTriangle(0);
	ray.setHitDistance(MAXFLOAT);

	stack.clear();
	
	return (Intersections::AABBRay(t_near, t_far, aabb.getMin(), aabb.getMax(),
							ray.getPoint(), ray.getInverseVector())) ? _itTraverse(root, ray,
					t_near, t_far) : MAXFLOAT;
}

#endif

